/**
   simulator.cu - Created by Timothy Morey on 2/2/2012.

   This file contains the implementation of the Simulator object.
*/

#include <curand_kernel.h>
#include <pthread.h>

#include "../nvn/nvn-common/client.h"
#include "../common/cuda.cuh"
#include "forest.cuh"
#include "nc_writer.h"
#include "simulator.h"

Simulator::Simulator(unsigned short width, unsigned short height)
  : _Forest(0),
		_ForestImg(0),
    _Width(width),
    _Height(height),
    _Running(false),
    _StopFlag(false),
    _ExportTarget(0)
{
  this->SetSize(width, height);
}

Simulator::~Simulator()
{
  this->Stop();
}

void Simulator::Start()
{
  _StopFlag = false;
  pthread_create(&_SimThread, 0, Simulator::SimThreadEntryPoint, this);
}

void Simulator::Stop()
{
  if(_Running)
  {
    _StopFlag = true;
    pthread_join(_SimThread, 0);
  }
}

void Simulator::SingleStep()
{
  if(this->IsRunning())
    this->Stop();

  unsigned int* curDevForest = 0;
  unsigned int* newDevForest = 0;
  size_t memsize = _Width * _Height * sizeof(unsigned int);
  CUDA_CALL( cudaMalloc(&curDevForest, memsize) );
  CUDA_CALL( cudaMalloc(&newDevForest, memsize) );
  CUDA_CALL( cudaMemcpy(curDevForest, _Forest, memsize, 
												cudaMemcpyHostToDevice) );
	
	char* devForestImg = 0;
	size_t imgsize = _Width * _Height * sizeof(char);
	CUDA_CALL( cudaMalloc(&devForestImg, imgsize) );
	CUDA_CALL( cudaMemcpy(devForestImg, _ForestImg, imgsize,
												cudaMemcpyHostToDevice) );

  curandState* rngStates;
  size_t rngsize = _Height * min(MAX_THREADS_PER_BLOCK, _Width) * sizeof(curandState);
  CUDA_CALL( cudaMalloc(&rngStates, rngsize) );
  
  InitRNG
    <<<_Height, min(MAX_THREADS_PER_BLOCK, _Width)>>>
    (rngStates, time(0));

  Step
    <<<_Height, min(MAX_THREADS_PER_BLOCK, _Width)>>>
    (curDevForest, newDevForest, rngStates, _Width, _Height);

	Downsample
		<<<_Height, min(MAX_THREADS_PER_BLOCK, _Width)>>>
		(newDevForest, devForestImg, _Width, _Height);
  
  CUDA_CALL( cudaMemcpy(_Forest, newDevForest, memsize, 
												cudaMemcpyDeviceToHost) );
	CUDA_CALL( cudaMemcpy(_ForestImg, devForestImg, imgsize,
												cudaMemcpyDeviceToHost) );

  if(_ExportTarget)
    Send(_ExportTarget, _ForestImg, imgsize);

  CUDA_CALL( cudaFree(curDevForest) );
  curDevForest = 0;

  CUDA_CALL( cudaFree(newDevForest) );
  newDevForest = 0;

	CUDA_CALL( cudaFree(devForestImg) );
	devForestImg = 0;
  
  CUDA_CALL( cudaFree(rngStates) );
  rngStates = 0;
}

void Simulator::SetSize(unsigned short width, unsigned short height)
{
  bool restart = false;
  if(this->IsRunning())
  {
    restart = true;
    this->Stop();
  }

  unsigned int* newForest = 0;
  size_t memsize = width * height * sizeof(unsigned int);
  CUDA_CALL( cudaHostAlloc(&newForest, memsize, cudaHostAllocDefault) );
  memset(newForest, 0, memsize);

	char* newForestImg = 0;
	size_t imgsize = width * height * sizeof(char);
	CUDA_CALL( cudaHostAlloc(&newForestImg, imgsize, cudaHostAllocDefault) );
	memset(newForestImg, 0, imgsize);

  if(_Forest)
  {
    CUDA_CALL( cudaFreeHost(_Forest) );
    _Forest = 0;
    _Width = 0;
    _Height = 0;
  }

	if(_ForestImg)
	{
		CUDA_CALL( cudaFreeHost(_ForestImg) );
		_ForestImg = 0;
	}

  _Forest = newForest;
	_ForestImg = newForestImg;
  _Width = width;
  _Height = height;

  if(restart)
    this->Start();
}

void Simulator::RandomizeForest()
{
  bool restart = false;
  if(this->IsRunning())
  {
    restart = true;
    this->Stop();
  }

  int nThreads = min(MAX_THREADS_PER_BLOCK, _Width);
  int nBlocks = _Height;

  unsigned int* devForest = 0;
  size_t memsize = _Width * _Height * sizeof(unsigned int);
  CUDA_CALL( cudaMalloc(&devForest, memsize) );
  CUDA_CALL( cudaMemcpy(devForest, _Forest, memsize, 
												cudaMemcpyHostToDevice) );

	char* devForestImg = 0;
	size_t imgsize = _Width * _Height * sizeof(char);
	CUDA_CALL( cudaMalloc(&devForestImg, imgsize) );
	CUDA_CALL( cudaMemcpy(devForestImg, _ForestImg, imgsize,
												cudaMemcpyHostToDevice) );

  curandState* rngStates = 0;
  size_t rngsize = nBlocks * nThreads * sizeof(curandState);
  CUDA_CALL( cudaMalloc(&rngStates, rngsize) );
  
  InitRNG 
    <<<nBlocks, nThreads>>> 
    (rngStates, time(0));

  // CUDA_CALL( cudaThreadSynchronize() );
  // CUDA_CALL( cudaGetLastError() );

  InitForest
    <<<nBlocks, nThreads>>>
    (devForest, rngStates, _Width, _Height);

	Downsample
		<<<nBlocks, nThreads>>>
		(devForest, devForestImg, _Width, _Height);

  // CUDA_CALL( cudaThreadSynchronize() );
  // CUDA_CALL( cudaGetLastError() );
  
  CUDA_CALL( cudaMemcpy(_Forest, devForest, memsize, 
												cudaMemcpyDeviceToHost) );
	CUDA_CALL( cudaMemcpy(_ForestImg, devForestImg, imgsize,
												cudaMemcpyDeviceToHost) );

  if(_ExportTarget)
    Send(_ExportTarget, _ForestImg, imgsize);

  CUDA_CALL( cudaFree(devForest) );
  devForest = 0;

	CUDA_CALL( cudaFree(devForestImg) );
	devForestImg = 0;

  CUDA_CALL( cudaFree(rngStates) );
  rngStates = 0;

  if(restart)
    this->Start();
}

void Simulator::PrintForest() const
{
  if(_Forest)
  {
    for(int i = 0; i < _Width * _Height; i++)
    {
      if(IsBurning(_Forest[i]))
				printf("X");
      else
				printf("%d", GetGrowth(_Forest[i]) / 26);
			
      if(_Width - 1 == (i % _Width))
				printf("\n");
    }
  }
}

void Simulator::ExportDataToClient(Client* client)
{
  if(_ExportTarget)
  {
		DestroyClient(_ExportTarget);
    _ExportTarget = 0;
  }

  _ExportTarget = client;

  if(_ExportTarget && _ForestImg)
    Send(_ExportTarget, _ForestImg, _Width * _Height * sizeof(char));
}

void* Simulator::SimThreadEntryPoint(void* arg)
{
  ((Simulator*)arg)->SimThreadLoop();
  return 0;
}

void Simulator::SimThreadLoop()
{
  _Running = true;

	int nblocks = _Height;
	int nthreads = min(MAX_THREADS_PER_BLOCK, _Width);

  unsigned int* curDevForest = 0;
  unsigned int* newDevForest = 0;
  size_t memsize = _Width * _Height * sizeof(unsigned int);
  CUDA_CALL( cudaMalloc(&curDevForest, memsize) );
  CUDA_CALL( cudaMalloc(&newDevForest, memsize) );
  CUDA_CALL( cudaMemcpy(curDevForest, _Forest, memsize, 
												cudaMemcpyHostToDevice) );
	
	char* devForestImg = 0;
	size_t imgsize = _Width * _Height * sizeof(char);
	CUDA_CALL( cudaMalloc(&devForestImg, imgsize) );
	CUDA_CALL( cudaMemcpy(devForestImg, _ForestImg, imgsize,
												cudaMemcpyHostToDevice) );
	
  curandState* rngStates;
  size_t rngsize = _Height * min(MAX_THREADS_PER_BLOCK, _Width) * sizeof(curandState);
  CUDA_CALL( cudaMalloc(&rngStates, rngsize) );
  
  InitRNG
    <<<nblocks, nthreads>>>
    (rngStates, time(0));

  while(! _StopFlag)
  {
    Step
      <<<nblocks, nthreads>>>
      (curDevForest, newDevForest, rngStates, _Width, _Height);

		Downsample
			<<<nblocks, nthreads>>>
			(newDevForest, devForestImg, _Width, _Height);

		CUDA_CALL( cudaMemcpy(_ForestImg, devForestImg, imgsize,
													cudaMemcpyDeviceToHost) );

    if(_ExportTarget)
      Send(_ExportTarget, _ForestImg, imgsize);

		curDevForest = newDevForest;
  }

	CUDA_CALL( cudaMemcpy(_Forest, curDevForest, memsize,
												cudaMemcpyDeviceToHost) );

  CUDA_CALL( cudaFree(curDevForest) );
  curDevForest = 0;

  CUDA_CALL( cudaFree(newDevForest) );
  newDevForest = 0;

	CUDA_CALL( cudaFree(devForestImg) );
	devForestImg = 0;
  
  CUDA_CALL( cudaFree(rngStates) );
  rngStates = 0;

  _Running = false;
}
